# 方向传播
"""
构建神经网路训练流程：
    1、准备数据集
    2、定义前馈计算 forward() 函数
    3、构造损失函数
    4、训练
        4.1、定义训练轮数 Epochs
            4.1.1、计算预测值 y'
            4.1.2、计算损失loss，这里会构建计算图，在这个计算图中会包括参数、输入、输出等信息，根据计算图进行前馈计算和反向传播
            4.1.3、方向传播 backward() 函数求梯度，这里会释放计算图，不会每次都生成一个计算图去占用资源
            4.1.4、更新参数，这里使用 w.grad.data，w.grad是一个张量，会生成计算图，浪费内存资源
            4.1.5、梯度清0（可选操作），防止梯度累加
        4.2、结束循环
    5、训练结束
"""

import torch

# 假设数据集
x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

# 设置参数，这里是权重
w = torch.Tensor([1.0])

# 设置自动求梯度
w.requires_grad = True


def forward(x):
    """
    前馈计算函数

    Args:
        x:输入值
    """
    return x * w


def loss(x, y):
    """
    计算损失

    Args:
         x:输入 x
         y:输入对应的实际值
    """
    # 计算 y'
    y_pred = forward(x)
    # 计算损失 loss
    return (y_pred - y) ** 2


# 这里 item() 返回是一个标量
print("predict (before training)", 4, forward(4).item())

# 迭代训练
for epoch in range(100):
    # 这里使用的是随机梯度下降
    for x, y in zip(x_data, y_data):
        # 计算损失
        l = loss(x, y)
        # 反向传播求梯度
        l.backward()
        # 输出梯度值
        print('\tgrad:', x, y, w.grad.item())
        # 更新参数，注意这里需要使用 data，因为 grad 是一个tensor，会构建计算图，有潜在问题，而且会浪费资源
        w.data = w.data - 0.01 * w.grad.data
        # 梯度清零，如果不清0，梯度梯度累加
        w.grad.data.zero_()
    print("progress:", epoch, l.item())

print("predict (after training)", 4, forward(4).item())
